# generate_answer.py

import argparse
import logging
import os
import re
import time
from typing import List, Dict
from .config import Config
import json
import requests


class AnswerGenerator:
    def __init__(self, config: Config):
        self.config = config

    @staticmethod
    def load_summaries(summary_file_path: str) -> List[Dict]:

        if not os.path.exists(summary_file_path):
            raise FileNotFoundError(f"找不到summary文件: {summary_file_path}")

        with open(summary_file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)

        return data.get("summarys", [])

    def find_relevant_document_prompt(self, question: str) -> str:

        try:
            summaries = self.load_summaries(self.config.summary)
            if not summaries:
                print("summary.json中没有找到文档总结")
                return "未找到文档总结"
        except Exception as e:
            print(f"加载文档总结时出错: {e}")
            return "加载文档总结时出错"
        summaries_text = "\n".join([
            f"文档: {summary['file_name']}\n总结: {summary['summary']}"
            for summary in summaries
        ])
        prompt = f"""
                你是一个智能文档分类助手。请根据用户提出的问题和提供的文档总结，判断该问题最可能属于哪个文档。

                问题: {question}

                文档总结列表:
                {summaries_text}

                要求:
                1. 仔细分析问题内容与各文档总结的相关性
                2. 选择最相关的几个文档
                3. 只返回文档名称，也就是文档中file_name字段的值，不要添加其他说明
                4. 如果最相关的文档有几个，且无法准确按断哪一个是最相关，则把最相关的几个文档全部返回
                5. 如果没有明显相关的文档，返回"未找到相关文档"
                """
        return prompt

    def find_most_relevant_document(self, question: str) -> str:

        prompt = self.find_relevant_document_prompt(question)
        logging.info(f"-----定位文档的prompt----{prompt}")
        # 调用大模型API
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json"
        }
        # Qwen3模型通过enable_thinking参数控制思考过程（开源版默认True，商业版默认False）
        # 使用Qwen3开源版模型时，若未启用流式输出，请将下行取消注释，否则会报错
        # extra_body={"enable_thinking": False},
        payload = {
            "model": self.config.model_name,
            "messages": [
                {"role": "system", "content": "你是一个专业的文档分类助手"},
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.3,  # 使用较低的temperature以获得更确定的结果
            "max_tokens": self.config.max_tokens,
            "enable_thinking": False
        }
        try:
            response = requests.post(
                self.config.api_endpoint,
                headers=headers,
                json=payload,
                timeout=self.config.timeout
            )
            response.raise_for_status()
            result = response.json()
            return result["choices"][0]["message"]["content"].strip()
        except Exception as e:
            print(f"调用大模型时出错: {e}")
            return "未找到相关文档"

    def find_document_in_markdown_folder(self, relevant_doc: str) -> str:
        """
        在markdown_folder及其子目录中查找文档
        """
        # 确保文件名以 .md 结尾
        if not relevant_doc.endswith('.md'):
            relevant_doc += '.md'
        # todo：大模型返回的文档名称可能不是严格遵守“文件夹/文件名.md”，这里要做各种兼容(返回的是 “文档：xxx”)，保证一定找到文件
        # todo：返回了多个文档名，要做循环处理
        for root, dirs, files in os.walk(self.config.markdown_folder):
            for file in files:
                file_path = os.path.join(root, file)
                relative_path = os.path.relpath(file_path, self.config.markdown_folder)
                # 如果 relevant_doc 包含斜杠，则按完整路径匹配
                if '/' in relevant_doc and relative_path == relevant_doc:
                    return file_path
                # 如果 relevant_doc 不包含斜杠，则仅比较文件名
                elif '/' not in relevant_doc and file == relevant_doc:
                    return file_path
        return ""

    def generate_answer_prompt(self, question: str, relevant_doc: str) -> str:
        # todo：未找到相关文档则遍历全部文档，并回答问题
        if not relevant_doc:
            return "未找到相关文档"

        markdown_file_path = self.find_document_in_markdown_folder(relevant_doc)
        if not markdown_file_path:
            return "markdown文档不存在"

        # 读取文档内容
        with open(markdown_file_path, 'r', encoding='utf-8') as f:
            doc_content = f.read()
        # todo: 调整提示词，给出json格式，判断文档和问题是否相关（用于遍历全部文档的场景）
        # 构建提示词
        prompt = f"""
        你是一个专业的文档问答助手。请基于以下文档内容，回答用户提出的问题。

        文档内容:
        {doc_content}

        问题:
        {question}

        要求:
        1. 回答必须完全基于文档内容，不要引入外部知识
        2. 如果文档中没有相关信息，请回答"文档中未提及此问题"
        3. 回答应简洁明了，直接返回答案内容，无需其他说明
        """
        return prompt

    def generate_answer(self, question: str) -> str:

        # 查找最相关的文档
        relevant_doc = self.find_most_relevant_document(question)
        think_tag = r'<think>.*?</think>'
        cleaned_think = re.sub(think_tag, '', relevant_doc, flags=re.DOTALL)
        relevant_doc = cleaned_think.strip()
        print(f"---relevant_doc----{relevant_doc}")
        prompt = self.generate_answer_prompt(question, relevant_doc)
        logging.info(f"-----生成答案的prompt----{prompt}")
        # 调用大模型API
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json"
        }

        payload = {
            "model": self.config.model_name,
            "messages": [
                {"role": "system", "content": "你是一个专业的文档问答助手"},
                {"role": "user", "content": prompt}
            ],
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "timeout": self.config.timeout,
            "enable_thinking": False
        }

        try:
            response = requests.post(
                self.config.api_endpoint,
                headers=headers,
                json=payload,
                timeout=self.config.timeout
            )
            response.raise_for_status()
            result = response.json()
            return result["choices"][0]["message"]["content"].strip()
        except Exception as e:
            return f"生成答案时出错: {e}"


def main():
    # 加载配置
    config = Config()

    # 解析命令行参数
    parser = argparse.ArgumentParser(description='根据问题和文档内容生成答案')
    parser.add_argument('question', type=str, help='需要回答的问题')
    args = parser.parse_args()

    # 创建答案生成器并生成答案
    generator = AnswerGenerator(config)
    start_time = time.time()
    answer = generator.generate_answer(args.question)

    # 计算并打印执行时间
    end_time = time.time()
    execution_time = end_time - start_time
    print(f"find_most_relevant_document 执行时间: {execution_time:.4f} 秒")
    print(answer)


if __name__ == "__main__":
    main()
